#include <array>
#include <cassert>
#include <cmath>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>

using namespace std;
// using namespace Eigen;
using arr4r = array<double, 4>;
using arr2r = array<double, 2>;
using arr3r = array<double, 3>;
using vecArr3r = vector<arr3r>;
using vecArr4r = vector<arr4r>;
using vecArr2r = vector<arr2r>;
template <typename T> using vec_t = std::vector<T>;

#include <Eigen/Dense>
#include <iostream>
/// 返回　x^order
inline double polyVal(const double x, const int order) {
  double val{1.0};
  for (int i = 1; i <= order; ++i)
    val *= x;
  return val;
}

void printUnorganizedPoint(std::ostream &os, const std::string &zone_name,
                           const arr2r *xyz, const int psize, double *rdata,
                           const std::string &data_name) {
  os << "VARIABLES = \"X\",\"Y\",\"";
  os << data_name << "\"\n";
  os << "ZONE T=\"" << zone_name << "\", I=" << psize
     << ", DATAPACKING=POINT\n";
  os << std::setprecision(16) << std::scientific;
  for (int i = 0; i < psize; ++i) {
    os << xyz[i][0] << " " << xyz[i][1];
    os << " " << rdata[i] << "\n";
  }
  os.flush();
}

void printGrid(const arr2r *xy, const int numi, const int numj,
               const double *sol, ostream &os, const string &zone_name,
               const std::string &data_name) {
  os << "VARIABLES = \"X\",\"Y\",\"";
  os << data_name << "\"\n";
  os << "ZONE T=\"" << zone_name << "\", I=" << numi << ", J=" << numj
     << ", DATAPACKING=POINT\n";
  os << std::setprecision(16) << std::scientific;
  for (int i = 0; i < numi; ++i)
    for (int j = 0; j < numj; ++j) {
      os << xy[i * numj + j][0] << "  " << xy[i * numj + j][1] << "  "
         << sol[i * numj + j] << "\n";
    }
  os.flush();
}

/**
 * @brief 测试多项式插值缩放的影响
 *
 */

/**
 * @brief 与上面的相比需要求出全部的系数
 * 不用求出逆了,因为需要所有的系数,后续可令一个单元分解一次,
 */
class WLSInterp {
public:
  int nrow_, ncol_, order_; ///< order为多项式最高次数
  Eigen::MatrixXd vand_;    ///< 可以少计算几次
  Eigen::MatrixXd wa_;      ///< vand*weights
  Eigen::ColPivHouseholderQR<Eigen::MatrixXd> qr_; ///< 分解有时可只做一次
  vec_t<double> s_;                                ///< 列平衡计算的系数

  /// 计算columnEquilibration()所需的缩放值
  void calcS() {
    for (int i = 0; i < ncol_; ++i)
      s_[i] = 1.0 / wa_.col(i).norm();
  }

public:
  Eigen::VectorXd poly_; ///< 给定插值点形成的代数多项式，系数均为1
  int ncol() const { return ncol_; }
  void init(int nrow, int poly_order) {
    nrow_ = nrow;
    order_ = poly_order;
    ncol_ = (order_ + 1) * (order_ + 2) / 2;
    vand_ = Eigen::MatrixXd::Zero(nrow_, ncol_);
    wa_ = Eigen::MatrixXd::Zero(nrow_, ncol_);
    s_.resize(ncol_);
    poly_ = Eigen::VectorXd(ncol_);
  }
  /// 构造范德蒙德矩阵的行
  void constructVandRow(const arr2r &dxy, const int row_idx) {
    vand_(row_idx, 0) = 1.0;
    int cnt{1};
    for (int od = 1; od <= order_; ++od)
      for (int i = 0; i < od + 1; ++i)
        vand_(row_idx, cnt++) = polyVal(dxy[0], od - i) * polyVal(dxy[1], i);
    assert(cnt == ncol_);
  }
  /// 已经计算出系数了,给出插值点的相对缩放坐标,给出多项式
  /// 对于一个单元可只计算一次
  void constructPoly(const arr2r &dxy, Eigen::VectorXd &poly) {
    assert(poly.size() == ncol_);
    poly(0) = 1.0;
    int cnt{1};
    for (int od = 1; od <= order_; ++od)
      for (int i = 0; i < od + 1; ++i)
        poly(cnt++) = polyVal(dxy[0], od - i) * polyVal(dxy[1], i);
  }

  void multiplyWeights(const double *w, const int num) {
    assert(num == nrow_);
    for (int i = 0; i < num; ++i)
      for (int j = 0; j < ncol_; ++j)
        wa_(i, j) = vand_(i, j) * w[i];
  }
  /// 估计不做这个也差不多
  void columnEquilibration() {
    calcS();
    for (int i = 0; i < ncol_; ++i)
      wa_.col(i) *= s_[i];
  }

  void invColumnEquilibration(double *coe, const int size) {
    for (int i = 0; i < ncol_; ++i)
      coe[i] *= s_[i];
  }
  /// 返回插值出的值
  double interpVal(double *coe, const int size) {
    assert(size == ncol_);
    double val{0};
    for (int i = 0; i < size; ++i)
      val += coe[i] * poly_(i);
    return val;
  }

  void colPivHouseholderQr() { qr_ = wa_.colPivHouseholderQr(); }

  void calcInterpCoe(const Eigen::VectorXd &b, double *coe,
                     const int coe_size) {
    assert(b.size() == nrow_ && coe_size == ncol_);
    Eigen::Map<Eigen::RowVectorXd> coef(coe, coe_size);
    coef = qr_.solve(b);
    // for (int i = 0; i < coe_size; ++i)
    //   coe[i] = f(i);
  }
};

struct Fun3 {
  double coef{1.0};
  double operator()(const arr2r &x) const {
    return (x[0] * x[0] * x[0] * x[0] * x[0] +
            x[1] * x[1] * x[1] * x[1] * x[1]) *
           coef;
    // return cos(x[0] * PI * 10) * sin(x[1] * PI * 10);
  }
};

int main(int argc, char *argv[]) {
  if (argc != 3) {
    std::cerr << "./a.out num dv" << std::endl;
    exit(EXIT_FAILURE);
  }
  const int num = stoi(argv[1]);
  const double dv = stod(argv[2]);
  double min_max[2]{0, 1};
  double dxdy[2]{1.0 / num, 2.0 / num};
  cout << setprecision(16) << scientific;

  Fun3 f;
  const int NUMP = 7, MID{NUMP / 2};
  arr2r xy[NUMP][NUMP];
  Eigen::VectorXd val(NUMP * NUMP);
  // 构造坐标
  int count{0};
  for (int i = 0; i < NUMP; ++i)
    for (int j = 0; j < NUMP; ++j) {
      xy[i][j][0] = i * dxdy[0];
      xy[i][j][1] = j * dxdy[1];
      if (j > MID)
        val(count++) = dv;
      else
        val(count++) = 0;
    }
  // 构造待插值点坐标
  const int per_ip = 5, size = NUMP * NUMP;
  arr2r ip[per_ip][per_ip];
  for (int i = 0; i < per_ip; ++i)
    for (int j = 0; j < per_ip; ++j) {
      ip[i][j][0] = xy[MID][MID][0] - 0.5 * dxdy[0] + dxdy[0] / per_ip * i;
      ip[i][j][1] = xy[MID][MID][1] - 0.5 * dxdy[1] + dxdy[1] / per_ip * j;
    }

  WLSInterp interp;
  interp.init(NUMP * NUMP, 2);
  double maxx{(NUMP - 1) * dxdy[0]}, maxy{(NUMP - 1) * dxdy[1]}, minx{0},
      miny{0}, centx{0.5 * (maxx + minx)}, centy{0.5 * (maxy + minx)};
  double rng = fmax(maxx - minx, maxy - miny);

  int cnt{0};
  for (int i = 0; i < NUMP; ++i)
    for (int j = 0; j < NUMP; ++j) {
      arr2r p{(xy[i][j][0] - centx) / rng, (xy[i][j][1] - centy) / rng};
      interp.constructVandRow(p, cnt);
      ++cnt;
    }
  const double h = dxdy[1], h2 = h * h, ep{1.0e-2};
  double weight[size], coe[interp.ncol()];
  Eigen::VectorXd b(size);
  const int mid_idx = MID * NUMP + MID;
  for (int i = 0; i < size; ++i) {
    weight[i] = (val[i] - val[mid_idx]) * (val[i] - val[mid_idx]) + ep * h2;
  }
  weight[mid_idx] = fmin(fmin(weight[mid_idx - 1], weight[mid_idx + 1]),
                         fmin(weight[mid_idx - NUMP], weight[mid_idx + NUMP]));
  for (int i = 0; i < size; ++i)
    weight[i] = 1.0 / weight[i];
  weight[mid_idx] *= 1.5;

  for (int i = 0; i < size; ++i) {
    // weight[i] = 1.0;
    cout << weight[i] << endl;
  }
  cout << endl;

  interp.multiplyWeights(weight, size);
  for (int i = 0; i < size; ++i)
    b[i] = weight[i] * val[i];
  // 缩放后值会很大,需要进行列平衡
  interp.columnEquilibration();
  interp.colPivHouseholderQr();
  interp.calcInterpCoe(b, coe, interp.ncol());
  interp.invColumnEquilibration(coe, interp.ncol());

  for (int i = 0; i < interp.ncol(); ++i)
    cout << coe[i] << endl;

  vec_t<double> ip_val(size);
  cnt = 0;
  for (int i = 0; i < per_ip; ++i)
    for (int j = 0; j < per_ip; ++j) {
      arr2r dxy;
      dxy[0] = (ip[i][j][0] - centx) / rng;
      dxy[1] = (ip[i][j][1] - centy) / rng;
      interp.constructPoly(dxy, interp.poly_);
      ip_val[cnt++] = interp.interpVal(coe, interp.ncol());
    }

  ofstream scatter{"scatter.plt"}, grid_val{"grid_val.plt"};

  printUnorganizedPoint(scatter, "scatter", &ip[0][0], per_ip * per_ip,
                        &ip_val[0], "val");
  printGrid(&xy[0][0], NUMP, NUMP, &val[0], grid_val, "grid_val", "val");
  scatter.close();
  grid_val.close();
  return 0;
}
